{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "519645be-a47b-4466-bbd2-16695a92d667",
   "metadata": {},
   "outputs": [],
   "source": [
    "## 导入必要库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "192394d7-db7f-4876-b812-4d7ff6aaf97a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from os import path as osp\n",
    "\n",
    "import hydra\n",
    "from omegaconf import DictConfig\n",
    "\n",
    "import ppsci\n",
    "from ppsci.utils import logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba1bdb0a-245a-4821-b8e2-96c5d78b554d",
   "metadata": {},
   "outputs": [],
   "source": [
    "## 定义模型训练模块\n",
    "初始化配置与环境设置：随机数种子、日志系统、是否启用并行计算\n",
    "~~~python\n",
    "ppsci.utils.misc.set_random_seed(cfg.seed)\n",
    "logger.init_logger(\"ppsci\", osp.join(cfg.output_dir, f\"{cfg.mode}.log\"), \"info\")\n",
    "~~~\n",
    "配置四个多层感知机（MLP），前两个为正向求解应变和应力的预训练模型，将参数冻结，后两个为分别求解λ和μ的模型\n",
    "~~~python\n",
    "disp_net = ppsci.arch.MLP(**cfg.MODEL.disp_net)\n",
    "stress_net = ppsci.arch.MLP(**cfg.MODEL.stress_net)\n",
    "inverse_lambda_net = ppsci.arch.MLP(**cfg.MODEL.inverse_lambda_net)\n",
    "inverse_mu_net = ppsci.arch.MLP(**cfg.MODEL.inverse_mu_net)\n",
    "~~~\n",
    "构建物理方程和几何约束\n",
    "~~~python\n",
    "equation = {\n",
    "        \"LinearElasticity\": ppsci.equation.LinearElasticity(\n",
    "            E=None, nu=None, lambda_=LAMBDA_, mu=MU, dim=3\n",
    "        )\n",
    "    }\n",
    "control_arm = ppsci.geometry.Mesh(cfg.GEOM_PATH)\n",
    "BOUNDS_X, BOUNDS_Y, BOUNDS_Z = control_arm.bounds\n",
    "~~~\n",
    "然后构造约束条件和验证器，最后初始化求解器并训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b580e033-750c-4277-8b52-383767191d19",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(cfg: DictConfig):\n",
    "    # set random seed for reproducibility\n",
    "    ppsci.utils.misc.set_random_seed(cfg.seed)\n",
    "    # initialize logger\n",
    "    logger.init_logger(\"ppsci\", osp.join(cfg.output_dir, f\"{cfg.mode}.log\"), \"info\")\n",
    "\n",
    "    # set model\n",
    "    disp_net = ppsci.arch.MLP(**cfg.MODEL.disp_net)\n",
    "    stress_net = ppsci.arch.MLP(**cfg.MODEL.stress_net)\n",
    "    inverse_lambda_net = ppsci.arch.MLP(**cfg.MODEL.inverse_lambda_net)\n",
    "    inverse_mu_net = ppsci.arch.MLP(**cfg.MODEL.inverse_mu_net)\n",
    "    # freeze models\n",
    "    disp_net.freeze()\n",
    "    stress_net.freeze()\n",
    "    # wrap to a model_list\n",
    "    model = ppsci.arch.ModelList(\n",
    "        (disp_net, stress_net, inverse_lambda_net, inverse_mu_net)\n",
    "    )\n",
    "\n",
    "    # set optimizer\n",
    "    lr_scheduler = ppsci.optimizer.lr_scheduler.ExponentialDecay(\n",
    "        **cfg.TRAIN.lr_scheduler\n",
    "    )()\n",
    "    optimizer = ppsci.optimizer.Adam(lr_scheduler)((inverse_lambda_net, inverse_mu_net))\n",
    "\n",
    "    # set equation\n",
    "    equation = {\n",
    "        \"LinearElasticity\": ppsci.equation.LinearElasticity(\n",
    "            E=None, nu=None, lambda_=\"lambda_\", mu=\"mu\", dim=3\n",
    "        )\n",
    "    }\n",
    "\n",
    "    # set geometry\n",
    "    control_arm = ppsci.geometry.Mesh(cfg.GEOM_PATH)\n",
    "    # geometry bool operation\n",
    "    geo = control_arm\n",
    "    geom = {\"geo\": geo}\n",
    "    # set bounds\n",
    "    BOUNDS_X, BOUNDS_Y, BOUNDS_Z = control_arm.bounds\n",
    "\n",
    "    # set dataloader config\n",
    "    interior_constraint = ppsci.constraint.InteriorConstraint(\n",
    "        equation[\"LinearElasticity\"].equations,\n",
    "        {\n",
    "            \"stress_disp_xx\": 0,\n",
    "            \"stress_disp_yy\": 0,\n",
    "            \"stress_disp_zz\": 0,\n",
    "            \"stress_disp_xy\": 0,\n",
    "            \"stress_disp_xz\": 0,\n",
    "            \"stress_disp_yz\": 0,\n",
    "        },\n",
    "        geom[\"geo\"],\n",
    "        {\n",
    "            \"dataset\": \"NamedArrayDataset\",\n",
    "            \"iters_per_epoch\": cfg.TRAIN.iters_per_epoch,\n",
    "            \"sampler\": {\n",
    "                \"name\": \"BatchSampler\",\n",
    "                \"drop_last\": True,\n",
    "                \"shuffle\": True,\n",
    "            },\n",
    "            \"num_workers\": 1,\n",
    "            \"batch_size\": cfg.TRAIN.batch_size.arm_interior,\n",
    "        },\n",
    "        ppsci.loss.MSELoss(\"sum\"),\n",
    "        criteria=lambda x, y, z: (\n",
    "            (BOUNDS_X[0] < x)\n",
    "            & (x < BOUNDS_X[1])\n",
    "            & (BOUNDS_Y[0] < y)\n",
    "            & (y < BOUNDS_Y[1])\n",
    "            & (BOUNDS_Z[0] < z)\n",
    "            & (z < BOUNDS_Z[1])\n",
    "        ),\n",
    "        name=\"INTERIOR\",\n",
    "    )\n",
    "    constraint = {interior_constraint.name: interior_constraint}\n",
    "\n",
    "    # set validator\n",
    "    LAMBDA_ = cfg.NU * cfg.E / ((1 + cfg.NU) * (1 - 2 * cfg.NU))  # 0.5769\n",
    "    MU = cfg.E / (2 * (1 + cfg.NU))  # 0.3846\n",
    "    geom_validator = ppsci.validate.GeometryValidator(\n",
    "        {\n",
    "            \"lambda_\": lambda out: out[\"lambda_\"],\n",
    "            \"mu\": lambda out: out[\"mu\"],\n",
    "        },\n",
    "        {\n",
    "            \"lambda_\": LAMBDA_,\n",
    "            \"mu\": MU,\n",
    "        },\n",
    "        geom[\"geo\"],\n",
    "        {\n",
    "            \"dataset\": \"NamedArrayDataset\",\n",
    "            \"sampler\": {\n",
    "                \"name\": \"BatchSampler\",\n",
    "                \"drop_last\": False,\n",
    "                \"shuffle\": False,\n",
    "            },\n",
    "            \"total_size\": cfg.EVAL.total_size.validator,\n",
    "            \"batch_size\": cfg.EVAL.batch_size.validator,\n",
    "        },\n",
    "        ppsci.loss.MSELoss(\"sum\"),\n",
    "        metric={\"L2Rel\": ppsci.metric.L2Rel()},\n",
    "        name=\"geo_eval\",\n",
    "    )\n",
    "    validator = {geom_validator.name: geom_validator}\n",
    "\n",
    "    # set visualizer(optional)\n",
    "    # add inferencer data\n",
    "    samples = geom[\"geo\"].sample_interior(\n",
    "        cfg.TRAIN.batch_size.visualizer_vtu,\n",
    "        criteria=lambda x, y, z: (\n",
    "            (BOUNDS_X[0] < x)\n",
    "            & (x < BOUNDS_X[1])\n",
    "            & (BOUNDS_Y[0] < y)\n",
    "            & (y < BOUNDS_Y[1])\n",
    "            & (BOUNDS_Z[0] < z)\n",
    "            & (z < BOUNDS_Z[1])\n",
    "        ),\n",
    "    )\n",
    "    pred_input_dict = {\n",
    "        k: v for k, v in samples.items() if k in cfg.MODEL.disp_net.input_keys\n",
    "    }\n",
    "    visualizer = {\n",
    "        \"visulzie_lambda_mu\": ppsci.visualize.VisualizerVtu(\n",
    "            pred_input_dict,\n",
    "            {\n",
    "                \"lambda\": lambda out: out[\"lambda_\"],\n",
    "                \"mu\": lambda out: out[\"mu\"],\n",
    "            },\n",
    "            prefix=\"vis\",\n",
    "        )\n",
    "    }\n",
    "\n",
    "    # initialize solver\n",
    "    solver = ppsci.solver.Solver(\n",
    "        model,\n",
    "        constraint,\n",
    "        cfg.output_dir,\n",
    "        optimizer,\n",
    "        lr_scheduler,\n",
    "        cfg.TRAIN.epochs,\n",
    "        cfg.TRAIN.iters_per_epoch,\n",
    "        seed=cfg.seed,\n",
    "        equation=equation,\n",
    "        geom=geom,\n",
    "        save_freq=cfg.TRAIN.save_freq,\n",
    "        log_freq=cfg.log_freq,\n",
    "        eval_freq=cfg.TRAIN.eval_freq,\n",
    "        eval_during_train=cfg.TRAIN.eval_during_train,\n",
    "        eval_with_no_grad=cfg.TRAIN.eval_with_no_grad,\n",
    "        validator=validator,\n",
    "        visualizer=visualizer,\n",
    "        pretrained_model_path=cfg.TRAIN.pretrained_model_path,\n",
    "    )\n",
    "\n",
    "    # train model\n",
    "    solver.train()\n",
    "\n",
    "    # plot losses\n",
    "    solver.plot_loss_history(by_epoch=False, smooth_step=1, use_semilogy=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8725c0fc-1627-45f8-a1b5-6d2811da931d",
   "metadata": {},
   "outputs": [],
   "source": [
    "## 定义模型评估模块"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "239ede66-60a5-40e5-adb2-effab25617a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(cfg: DictConfig):\n",
    "    # set random seed for reproducibility\n",
    "    ppsci.utils.misc.set_random_seed(cfg.seed)\n",
    "    # initialize logger\n",
    "    logger.init_logger(\"ppsci\", osp.join(cfg.output_dir, f\"{cfg.mode}.log\"), \"info\")\n",
    "\n",
    "    # set model\n",
    "    disp_net = ppsci.arch.MLP(**cfg.MODEL.disp_net)\n",
    "    stress_net = ppsci.arch.MLP(**cfg.MODEL.stress_net)\n",
    "    inverse_lambda_net = ppsci.arch.MLP(**cfg.MODEL.inverse_lambda_net)\n",
    "    inverse_mu_net = ppsci.arch.MLP(**cfg.MODEL.inverse_mu_net)\n",
    "    # wrap to a model_list\n",
    "    model = ppsci.arch.ModelList(\n",
    "        (disp_net, stress_net, inverse_lambda_net, inverse_mu_net)\n",
    "    )\n",
    "\n",
    "    # set geometry\n",
    "    control_arm = ppsci.geometry.Mesh(cfg.GEOM_PATH)\n",
    "    # geometry bool operation\n",
    "    geo = control_arm\n",
    "    geom = {\"geo\": geo}\n",
    "    # set bounds\n",
    "    BOUNDS_X, BOUNDS_Y, BOUNDS_Z = control_arm.bounds\n",
    "\n",
    "    # set validator\n",
    "    LAMBDA_ = cfg.NU * cfg.E / ((1 + cfg.NU) * (1 - 2 * cfg.NU))  # 0.57692\n",
    "    MU = cfg.E / (2 * (1 + cfg.NU))  # 0.38462\n",
    "    geom_validator = ppsci.validate.GeometryValidator(\n",
    "        {\n",
    "            \"lambda_\": lambda out: out[\"lambda_\"],\n",
    "            \"mu\": lambda out: out[\"mu\"],\n",
    "        },\n",
    "        {\n",
    "            \"lambda_\": LAMBDA_,\n",
    "            \"mu\": MU,\n",
    "        },\n",
    "        geom[\"geo\"],\n",
    "        {\n",
    "            \"dataset\": \"NamedArrayDataset\",\n",
    "            \"sampler\": {\n",
    "                \"name\": \"BatchSampler\",\n",
    "                \"drop_last\": False,\n",
    "                \"shuffle\": False,\n",
    "            },\n",
    "            \"total_size\": cfg.EVAL.total_size.validator,\n",
    "            \"batch_size\": cfg.EVAL.batch_size.validator,\n",
    "        },\n",
    "        ppsci.loss.MSELoss(\"sum\"),\n",
    "        metric={\"L2Rel\": ppsci.metric.L2Rel()},\n",
    "        name=\"geo_eval\",\n",
    "    )\n",
    "    validator = {geom_validator.name: geom_validator}\n",
    "\n",
    "    # set visualizer(optional)\n",
    "    # add inferencer data\n",
    "    samples = geom[\"geo\"].sample_interior(\n",
    "        cfg.EVAL.batch_size.visualizer_vtu,\n",
    "        criteria=lambda x, y, z: (\n",
    "            (BOUNDS_X[0] < x)\n",
    "            & (x < BOUNDS_X[1])\n",
    "            & (BOUNDS_Y[0] < y)\n",
    "            & (y < BOUNDS_Y[1])\n",
    "            & (BOUNDS_Z[0] < z)\n",
    "            & (z < BOUNDS_Z[1])\n",
    "        ),\n",
    "    )\n",
    "    pred_input_dict = {\n",
    "        k: v for k, v in samples.items() if k in cfg.MODEL.disp_net.input_keys\n",
    "    }\n",
    "    visualizer = {\n",
    "        \"visulzie_lambda_mu\": ppsci.visualize.VisualizerVtu(\n",
    "            pred_input_dict,\n",
    "            {\n",
    "                \"lambda\": lambda out: out[\"lambda_\"],\n",
    "                \"mu\": lambda out: out[\"mu\"],\n",
    "            },\n",
    "            prefix=\"vis\",\n",
    "        )\n",
    "    }\n",
    "\n",
    "    # initialize solver\n",
    "    solver = ppsci.solver.Solver(\n",
    "        model,\n",
    "        output_dir=cfg.output_dir,\n",
    "        seed=cfg.seed,\n",
    "        log_freq=cfg.log_freq,\n",
    "        eval_with_no_grad=cfg.EVAL.eval_with_no_grad,\n",
    "        validator=validator,\n",
    "        visualizer=visualizer,\n",
    "        pretrained_model_path=cfg.EVAL.pretrained_model_path,\n",
    "    )\n",
    "    # evaluate after finished training\n",
    "    solver.eval()\n",
    "    # visualize prediction after finished training\n",
    "    solver.visualize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22b964fe-26ed-4a79-90f4-a52b562b887f",
   "metadata": {},
   "outputs": [],
   "source": [
    "## 定义模型导出模块"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "703e8df6-3214-4dd5-bc30-7bb91d944f3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def export(cfg: DictConfig):\n",
    "    from paddle.static import InputSpec\n",
    "\n",
    "    # set model\n",
    "    disp_net = ppsci.arch.MLP(**cfg.MODEL.disp_net)\n",
    "    stress_net = ppsci.arch.MLP(**cfg.MODEL.stress_net)\n",
    "    inverse_lambda_net = ppsci.arch.MLP(**cfg.MODEL.inverse_lambda_net)\n",
    "    inverse_mu_net = ppsci.arch.MLP(**cfg.MODEL.inverse_mu_net)\n",
    "    # wrap to a model_list\n",
    "    model = ppsci.arch.ModelList(\n",
    "        (disp_net, stress_net, inverse_lambda_net, inverse_mu_net)\n",
    "    )\n",
    "\n",
    "    # load pretrained model\n",
    "    solver = ppsci.solver.Solver(\n",
    "        model=model, pretrained_model_path=cfg.INFER.pretrained_model_path\n",
    "    )\n",
    "\n",
    "    # export models\n",
    "    input_spec = [\n",
    "        {\n",
    "            key: InputSpec([None, 1], \"float32\", name=key)\n",
    "            for key in cfg.MODEL.disp_net.input_keys\n",
    "        },\n",
    "    ]\n",
    "    solver.export(input_spec, cfg.INFER.export_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f278c273-c6c4-4c55-8a1d-a1a3a359916e",
   "metadata": {},
   "outputs": [],
   "source": [
    "## 定义模型推理模块"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68e87509-3303-4016-9656-575f57cae575",
   "metadata": {},
   "outputs": [],
   "source": [
    "def inference(cfg: DictConfig):\n",
    "    from deploy.python_infer import pinn_predictor\n",
    "    from ppsci.visualize import vtu\n",
    "\n",
    "    # set model predictor\n",
    "    predictor = pinn_predictor.PINNPredictor(cfg)\n",
    "\n",
    "    # set geometry\n",
    "    control_arm = ppsci.geometry.Mesh(cfg.GEOM_PATH)\n",
    "    # geometry bool operation\n",
    "    geo = control_arm\n",
    "    geom = {\"geo\": geo}\n",
    "    # set bounds\n",
    "    BOUNDS_X, BOUNDS_Y, BOUNDS_Z = control_arm.bounds\n",
    "    samples = geom[\"geo\"].sample_interior(\n",
    "        cfg.EVAL.batch_size.visualizer_vtu,\n",
    "        criteria=lambda x, y, z: (\n",
    "            (BOUNDS_X[0] < x)\n",
    "            & (x < BOUNDS_X[1])\n",
    "            & (BOUNDS_Y[0] < y)\n",
    "            & (y < BOUNDS_Y[1])\n",
    "            & (BOUNDS_Z[0] < z)\n",
    "            & (z < BOUNDS_Z[1])\n",
    "        ),\n",
    "    )\n",
    "    pred_input_dict = {\n",
    "        k: v for k, v in samples.items() if k in cfg.MODEL.disp_net.input_keys\n",
    "    }\n",
    "\n",
    "    output_dict = predictor.predict(pred_input_dict, cfg.INFER.batch_size)\n",
    "\n",
    "    # mapping data to output_keys\n",
    "    output_keys = (\n",
    "        cfg.MODEL.disp_net.output_keys\n",
    "        + cfg.MODEL.stress_net.output_keys\n",
    "        + cfg.MODEL.inverse_lambda_net.output_keys\n",
    "        + cfg.MODEL.inverse_mu_net.output_keys\n",
    "    )\n",
    "    output_dict = {\n",
    "        store_key: output_dict[infer_key]\n",
    "        for store_key, infer_key in zip(output_keys, output_dict.keys())\n",
    "    }\n",
    "    output_dict.update(pred_input_dict)\n",
    "    vtu.save_vtu_from_dict(\n",
    "        osp.join(cfg.output_dir, \"vis\"),\n",
    "        output_dict,\n",
    "        cfg.MODEL.disp_net.input_keys,\n",
    "        output_keys,\n",
    "        1,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4461b925-4c35-47cd-9a94-56a90353082c",
   "metadata": {},
   "outputs": [],
   "source": [
    "## 主函数入口"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08b3df08-7b2c-4b0b-bbd4-fbf1a358ea28",
   "metadata": {},
   "outputs": [],
   "source": [
    "@hydra.main(\n",
    "    version_base=None, config_path=\"../conf\", config_name=\"inverse_parameter.yaml\"\n",
    ")\n",
    "def main(cfg: DictConfig):\n",
    "    if cfg.mode == \"train\":\n",
    "        train(cfg)\n",
    "    elif cfg.mode == \"eval\":\n",
    "        evaluate(cfg)\n",
    "    elif cfg.mode == \"export\":\n",
    "        export(cfg)\n",
    "    elif cfg.mode == \"infer\":\n",
    "        inference(cfg)\n",
    "    else:\n",
    "        raise ValueError(\n",
    "            f\"cfg.mode should in ['train', 'eval', 'export', 'infer'], but got '{cfg.mode}'\"\n",
    "        )\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
